use libc::*;
use libnuma_sys::*;
use std::ffi::{CStr, CString, NulError};
use std::mem::MaybeUninit;
use std::num::ParseIntError;
use std::ptr::{null, null_mut};
use std::str::FromStr;
use std::time::Duration;
fn main() {
    let mut cpu_core: Vec<usize> = Vec::new();
    let mut mem_node: Vec<c_uint> = Vec::new();
    let mut exe_path: Vec<String> = Vec::new();
    std::env::args().skip(1).for_each(|arg| {
        if arg.starts_with("--cpubind=") {
            let mut vec = resolve_args::<usize>(arg.clone(), "--cpubind=");
            cpu_core.append(&mut vec);
        } else if arg.starts_with("--membind=") {
            let mut vec = resolve_args::<c_uint>(arg.clone(), "--membind=");
            mem_node.append(&mut vec);
        } else {
            exe_path.push(arg.clone());
        }
    });
    if cpu_core.len() == 0 && mem_node.len() == 0 {
        print_banner();
        unsafe {
            exit(1);
        }
    }

    unsafe {
        let i = fork();
        if i != 0 {
            println!("[info]: create new process pid:{}", i);
        }
        if i == 0 {
            setting_numa(
                {
                    if cpu_core.len() > 0 {
                        Some(cpu_core)
                    } else {
                        None
                    }
                },
                {
                    if mem_node.len() > 0 {
                        Some(mem_node)
                    } else {
                        None
                    }
                },
            );
            let mut exe_path_cstring = exe_path
                .iter()
                .map(|s| {
                    let cstr = match CString::new(s.as_str()) {
                        Ok(cstr) => {
                            cstr},
                        Err(_) => exit(1),
                    };
                    cstr
                })
                .collect::<Vec<_>>();
            let exe_path_cstr = exe_path_cstring.iter().map(|c| { c.as_ptr() }).collect::<Vec<_>>();
            let i1 = execvp(exe_path_cstr[0],exe_path_cstr.as_ptr() );
            let errno_location = __errno_location();
            printf(c"[error]: %s\n".as_ptr(),strerror(*errno_location));
            println!("[error]: error create process",);
            println!("{:?}",exe_path)
        }
    }
}
fn resolve_args<T>(args: String, prefix: &str) -> Vec<T>
where
    T: FromStr<Err = ParseIntError>,
{
    let mut args_n = match args.strip_prefix(prefix) {
        None => {
            print_banner();
            unsafe {
                exit(1);
            }
        }
        Some(nums) => nums
            .split(',')
            .map(|mut n_str| match n_str.trim().parse::<T>() {
                Ok(n) => n,
                Err(e) => {
                    print_banner();
                    unsafe {
                        exit(1);
                    }
                }
            })
            .collect::<Vec<T>>(),
    };
    args_n
}
fn print_banner() {
    println!();
    println!("usage:");
    println!("<this utility> [--cpubind=<0,[1[,2[,3...]>] [--membind=<0,[1[,2[,3...]> <program_executable>]");
    println!();
    println!("exmple:");
    println!("<this utility> --cpubind=0,1,2 --membind=0 program");
    println!();
}

fn setting_numa(cpu_core: Option<Vec<usize>>, mem_node: Option<Vec<c_uint>>) {
    unsafe {
        println!("[info]: check numa available");
        if numa_available() == -1 {
            println!("[error]: numa库不可用");
            exit(1);
        }
        if let Some(cpu_core) = cpu_core {
            let mut cpu_set_t = MaybeUninit::<cpu_set_t>::uninit().assume_init();
            CPU_ZERO(&mut cpu_set_t);
            //设置要启用的cpu掩码
            cpu_core.into_iter().for_each(|cpu_core| {
                CPU_SET(cpu_core, &mut cpu_set_t);
            });
            //开始起用
            println!("[info]: set process use numa cpu core");
            if sched_setaffinity(getpid(), size_of::<cpu_set_t>(), &cpu_set_t) == -1 {
                println!("[error]: sched_setaffinity failed");
                println!("[error]: cpu核心不可设置");
                exit(1);
            };
        }
        if let Some(mem_node) = mem_node {
            let mut bitmask_alloc = numa_bitmask_alloc((numa_max_node() + 1) as c_uint);
            mem_node.into_iter().for_each(|mem_node| {
                numa_bitmask_setbit(bitmask_alloc, mem_node);
            });
            // numa_set_membind(membind_nodes, NUMA_MEMBIND_BIND, MPOL_MF_STRICT);
            println!("[info]: allocate memnode");
            numa_set_membind(bitmask_alloc);
            // const size_t size = 1024;
            // void *ptr = numa_alloc(size);  // 使用NUMA-aware内存分配
            // if (!ptr) {
            //     perror("numa_alloc失败");
            //     return 1;
            // }
            let test_alloc = numa_alloc(1024);
            if test_alloc == null_mut() {
                println!("[error]: test_alloc failed");
                println!("[error]: numa 内存节点分配失败");
                exit(1);
            }
            numa_free(test_alloc, 1024);
            numa_bitmask_free(bitmask_alloc);
        }
    }
}
